package org.laizili.solution.leetcode;

import java.util.*;

/**
 * <a href="https://leetcode-cn.com/problems/random-pick-index/">398. 随机数索引</a>
 * <p>
 * tags: 水塘抽样;
 */
public class Problem398 {
    private static class Solution {
        Map<Integer, List<Integer>> mapper = new HashMap<>();
        Random random = new Random();
        int[] numbers;

        public Solution(int[] nums) {
            for (int i = 0; i < nums.length; i++) {
                mapper.putIfAbsent(nums[i], new ArrayList<>());
                mapper.get(nums[i]).add(i);
            }
            numbers = nums;
        }

        public int pick(int target) {
            List<Integer> li = mapper.get(target);
            return li.get(random.nextInt(li.size()));
        }
    }

    // 水塘抽样
    private static class Solution2 {
        private final int[] nums;
        Random random = new Random();

        public Solution2(int[] nums) {
            this.nums = nums;
        }

        public int pick(int target) {
            int i = 0;
            int counter = 0;
            for (int k = 0; k < nums.length; k++) {
                if (nums[k] == target) {
                    counter++;
                    if (random.nextInt(counter) == 0) {
                        i = k;
                    }
                }
            }
            return i;
        }
    }

    public static void main(String[] args) {
        final Solution2 solution = new Solution2(new int[]{1, 2, 3, 3, 3});
        System.out.println(solution.pick(3));
        System.out.println(solution.pick(1));
        System.out.println(solution.pick(3));

    }
}
